from enum import StrEnum
from datasets import load_from_disk
from torch.utils.data import DataLoader
from configuration import config


class DataType(StrEnum):
    TRAIN = "train"
    TEST = "test"
    VALID = "valid"


def get_dataset(type):
    dataset = load_from_disk(str(config.PROCESSED_DATA_DIR / type))
    dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
    return dataset


def get_dataloader(type=DataType.TRAIN):
    dataset = get_dataset(type)
    return DataLoader(dataset, batch_size=config.BATCH_SIZE, shuffle=True)
